package br.org.game.nio.core.tcp;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.net.SocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.ClosedChannelException;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;

public class SocketServer {

	private static final int SERVER_PORT = 1234;

	private ServerSocketChannel serverSocketChannel;
	private Selector selector;

	private Map<SocketChannel, byte[]> clientsMap;

	public SocketServer() throws IOException {
		clientsMap = new HashMap<SocketChannel, byte[]>();

		configureServer();

		startServer();
	}

	private void configureServer() throws IOException, ClosedChannelException {
		serverSocketChannel = ServerSocketChannel.open();

		serverSocketChannel.configureBlocking(false); // no blocking socket

		serverSocketChannel.socket().bind(new InetSocketAddress("localhost", SERVER_PORT)); // bind server to address

		selector = Selector.open();

		serverSocketChannel.register(selector, SelectionKey.OP_ACCEPT); // registering the selector in socket server
	}

	private void startServer() throws IOException {

		while (true) {

			selector.select(); // wait for events: accept, read and write.

			Iterator<SelectionKey> iteratorKeys = selector.selectedKeys().iterator(); // selected keys

			while (iteratorKeys.hasNext()) {

				SelectionKey key = iteratorKeys.next();

				iteratorKeys.remove(); // this is necessary to prevent the same key from coming up again the next time around.

				if (!key.isValid()) {

					continue;

				}

				processKey(key); // process event
			}

		}

	}

	private void processKey(SelectionKey key) throws IOException {
		if (key.isAcceptable()) {

			this.accept(key);

		}

		if (key.isValid() && key.isReadable()) {

			this.read(key);

		}

		if (key.isValid() && key.isWritable()) {

			this.write(key);

		}
	}

	private void accept(SelectionKey key) throws IOException {
		ServerSocketChannel serverChannel = (ServerSocketChannel) key.channel();

		SocketChannel channel = serverChannel.accept();

		channel.configureBlocking(false);

		Socket socket = channel.socket();

		SocketAddress remoteAddr = socket.getRemoteSocketAddress();

		System.out.println("Connected to: " + remoteAddr);

		clientsMap.put(channel, null);

		channel.register(selector, SelectionKey.OP_READ); // register channel with selector for further read
	}

	private void read(SelectionKey key) throws IOException {
		SocketChannel channel = (SocketChannel) key.channel();

		ByteBuffer buffer = ByteBuffer.allocate(29);

		int numRead = channel.read(buffer); // read message

		if (numRead <= 0) { // if end of stream channel

			closeChannel(key); // close channel, cancel key and remove from storage

			return;
		}

		byte[] data = new byte[numRead];

		System.arraycopy(buffer.array(), 0, data, 0, numRead);

		System.out.println("READ: " + new String(data));

		doEcho(key, data); // write back to all client, inclusive yourself
	}

	private void write(SelectionKey key) throws IOException {
		SocketChannel channel = (SocketChannel) key.channel();

		byte[] pendingData = clientsMap.get(channel); // get stored data

		if (pendingData != null) {

			Set<SocketChannel> channels = clientsMap.keySet(); // get all connected clients for write

			for (SocketChannel socketChannel : channels) {

				socketChannel.write(ByteBuffer.wrap(pendingData)); // write in the channel

				System.out.println("WRITE: " + new String(pendingData));

			}

		}

		clientsMap.put(channel, null); // remove writen data

		key.interestOps(SelectionKey.OP_READ); // register channel with selector for further read
	}

	private void closeChannel(SelectionKey key) throws IOException {
		SocketChannel channel = (SocketChannel) key.channel();

		clientsMap.remove(channel); // remove from storage

		Socket socket = channel.socket();

		SocketAddress remoteAddr = socket.getRemoteSocketAddress();

		System.out.println("Connection closed by client: " + remoteAddr);

		channel.close(); // close channel

		key.cancel(); // cancel key for next operations
	}

	private void doEcho(SelectionKey key, byte[] data) {
		SocketChannel channel = (SocketChannel) key.channel();

		clientsMap.put(channel, data); // store data read

		key.interestOps(SelectionKey.OP_WRITE); // register channel with selector for further write
	}

	public static void main(String[] args) throws IOException {

		new SocketServer();

	}

}